# label_shift_estimator.py
# 使用EM算法实现 MLLS (Maximum Likelihood Label Shift) 估计。
# 复用您提供的代码。

import torch
import numpy as np

class MLLSEstimator:
    """
    使用期望最大化 (EM) 算法来估计标签偏移。
    """
    def __init__(self, num_classes, max_iter=1000, tol=1e-6):
        self.num_classes = num_classes
        self.max_iter = max_iter
        self.tol = tol
        
    def estimate_target_priors(self, target_probs, source_priors):
        """
        使用EM算法估计目标域的先验 p_t(y)。
        返回:
            torch.Tensor: 估计出的目标域先验概率向量 p_t(y)。
        """
        # 数值稳定性检查
        eps = 1e-8
        source_priors = torch.clamp(source_priors, min=eps)
        target_probs = torch.clamp(target_probs, min=eps)
        
        # 归一化确保概率和为1
        source_priors = source_priors / source_priors.sum()
        target_probs = target_probs / target_probs.sum(dim=1, keepdim=True)
        
        # 初始化 p_t(y) 为源域先验
        p_t_y = source_priors.clone()
        
        for i in range(self.max_iter):
            p_t_y_old = p_t_y.clone()
            
            # E-步 和 M-步 结合
            w = p_t_y / (source_priors + eps)
            
            # 限制权重避免过大值
            w = torch.clamp(w, min=eps, max=100.0)
            
            # p_hat(x) = sum_y p_hat(y|x) * w(y) * p_s(y)
            p_x_est = target_probs @ (w * source_priors)
            p_x_est = torch.clamp(p_x_est, min=eps)

            # p_hat(y|x) = (p_hat(y|x) * w(y) * p_s(y)) / p_hat(x)
            posteriors = (target_probs * (w * source_priors)) / p_x_est.unsqueeze(1)
            
            # M-步: 更新 p_t(y) 的估计
            p_t_y = torch.mean(posteriors, dim=0)
            
            # 确保概率有效性
            p_t_y = torch.clamp(p_t_y, min=eps)
            p_t_y = p_t_y / p_t_y.sum()  # 归一化
            
            # 检查收敛
            if torch.norm(p_t_y - p_t_y_old) < self.tol:
                print(f"EM algorithm converged at iteration {i+1}")
                break
        
        if i == self.max_iter - 1:
            print("Warning: EM algorithm did not converge within max_iter.")
            
        return p_t_y
